# MySQL数据库操作辅助类
import mysql.connector
from mysql.connector import Error
import logging
from app.conf.config import config

class MySQLHelper:
    def __init__(self, auto_commit=False):
        self.db_config = {
            'host': config.DB_HOST,
            'user': config.DB_USER,
            'password': config.DB_PASSWORD,
            'database': config.DB_NAME
        }
        self.logger = logging.getLogger(__name__)
        self.conn = None
        self.auto_commit = auto_commit
        self.in_transaction = False
    
    def _get_connection(self):
        """获取数据库连接，保持连接重用"""
        try:
            # 检查连接是否存在且有效
            if self.conn and self.conn.is_connected():
                return self.conn
            
            # 创建新连接
            self.conn = mysql.connector.connect(
                host=self.db_config['host'],
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database'],
                autocommit=self.auto_commit
            )
            
            return self.conn
        except Error as e:
            self.logger.error(f"数据库连接失败: {e}")
            # 发生错误时重置连接
            self.conn = None
            raise
    
    def close(self):
        """显式关闭连接"""
        if self.conn and self.conn.is_connected():
            try:
                self.conn.close()
                self.logger.info("数据库连接已关闭")
                self.conn = None
            except Error as e:
                self.logger.error(f"关闭连接时出错: {e}")
    
    def execute_query(self, query, params=None):
        """执行SQL查询，支持连接重用"""
        connection = None
        cursor = None
        
        try:
            connection = self._get_connection()
            cursor = connection.cursor(dictionary=True)
            cursor.execute(query, params or ())
            result = cursor.fetchall()
            return result
        except Error as e:
            self.logger.error(f"执行查询失败: {e}")
            # 如果不在事务中且发生错误，重置连接
            if not self.in_transaction:
                self.conn = None
            raise
        finally:
            # 只关闭游标，不关闭连接
            if cursor:
                try:
                    cursor.close()
                except Error as e:
                    self.logger.warning(f"关闭游标时出错: {e}")
    
    def execute_update(self, query, params=None):
        """执行更新操作，支持连接重用"""
        connection = None
        cursor = None
        
        try:
            connection = self._get_connection()
            cursor = connection.cursor()
            cursor.execute(query, params or ())
            
            # 如果不在事务中且设置了自动提交，则提交
            if not self.in_transaction and self.auto_commit:
                connection.commit()
            
            return True
        except Error as e:
            self.logger.error(f"执行更新失败: {e}")
            # 如果不在事务中，尝试回滚并重置连接
            if not self.in_transaction:
                try:
                    if connection:
                        connection.rollback()
                except:
                    pass
                self.conn = None
            raise
        finally:
            # 只关闭游标，不关闭连接
            if cursor:
                try:
                    cursor.close()
                except Error as e:
                    self.logger.warning(f"关闭游标时出错: {e}")
    
    def transaction(self):
        """事务上下文管理器"""
        return TransactionContext(self)


class TransactionContext:
    """事务上下文管理器，用于处理数据库事务"""
    
    def __init__(self, mysql_helper):
        self.mysql_helper = mysql_helper
        self.logger = logging.getLogger(__name__)
    
    def __enter__(self):
        """开始事务"""
        try:
            conn = self.mysql_helper._get_connection()
            self.mysql_helper.in_transaction = True
            conn.autocommit = False  # 确保事务模式
            self.logger.info("事务开始")
            return conn
        except Exception as e:
            self.logger.error(f"开始事务失败: {e}")
            raise
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """结束事务，根据异常情况决定提交或回滚"""
        try:
            conn = self.mysql_helper.conn
            if conn and conn.is_connected():
                if exc_type is None:
                    # 没有异常，提交事务
                    conn.commit()
                    self.logger.info("事务提交成功")
                else:
                    # 有异常，回滚事务
                    try:
                        conn.rollback()
                        self.logger.info("事务回滚成功")
                    except Exception as rollback_error:
                        self.logger.error(f"事务回滚失败: {rollback_error}")
        finally:
            # 无论如何，重置事务状态
            self.mysql_helper.in_transaction = False
            # 如果设置了自动提交，恢复自动提交模式
            if self.mysql_helper.auto_commit and self.mysql_helper.conn and self.mysql_helper.conn.is_connected():
                self.mysql_helper.conn.autocommit = True